# Make coefficient plots 

library(tidyverse)

# DEFINE COLORS FOR CONSISTENCY ----
# create a named vector for consistent colors across plots with different treatment rules
color_palette <- scales::hue_pal()(3)
names(color_palette) <- c("one-dimensional", "quadrant", "cubic")
color_palette <- color_palette[c(2, 3, 1)] # reorder colors


##############################################################################################
#                                                                                            #          
#                  Percent to treat as text labels in coefficient plots                      # 
#                                                                                            # 
##############################################################################################

# MAIN PLOTS: DEMOGRAPHICS AND PRE-TREATMENT ELECTRICITY CONSUMPTION ----
make_separate_plot_text <- function(filename) {
  raw.dat <- readLines(paste0(filename,".txt"), warn = FALSE)
  
  tab.dat <- raw.dat %>% 
    gsub("\\\\|\\(|)| ", "", .) %>% 
    read.table(text = ., sep = "&", stringsAsFactors = FALSE)
  
  if(tab.dat[1,1]=="ActualRCT"){
    tab.dat <- tab.dat[-(1:2),]  
    tab.dat <- tab.dat[,-(4)] 
  }
  
  dat <- cbind(tab.dat[grepl("EWM",tab.dat[,1]), ], tab.dat[!grepl("EWM",tab.dat[,1]), 4])
  names(dat) <- c("type", "variable", 'share',"pe", "ci") 
  dat <- dat %>% separate(ci, c("lb","ub"), ",")
  dat[,3:ncol(dat)] <- sapply(dat[,3:ncol(dat)], as.numeric)
  dat$type <- gsub("EWM-", "", dat$type)
  dat$type <- str_to_lower(dat$type)
  dat <- dat %>% mutate(variable = ifelse(variable=="std", "sd", variable),
                        variable = ifelse(variable=="size", "house size", variable),
                        variable = ifelse(variable=="vintage", "house age", variable),
                        variable = paste0(variable," and \n mean electricity usage"),
                        variable = ifelse(type=="univariate" & grepl("income", variable), "only \n mean electricity usage", variable))

  if(gsub("^.*_", "", filename)=="kwh"){
    ylabel <- "Electricity Consumption Changes \n (kWh per household per month)"
    limits <- c(-7.5,2) 
    breaks <- seq(2,-7,-2)
  } else{
    ylabel <- "Net Cost Changes \n ($ per household per month)"
    limits <- c(-1.5, 0) 
    breaks <- seq(0.5,-1.25,-0.25)
  }
  
  if (max(dat$ub)<0){ 
    y_text <- -1.5
  } else{
    y_text <- -7 
  }
  
  if (gsub("_.*", "", filename)=="demographics"){ 
    ewm_rules <- c("quadrant", "cubic")
    ewm_colors <- color_palette[c("quadrant", "cubic")]
  } else{
     ewm_rules <- c("quadrant", "cubic", "one-dimensional")
    ewm_colors <- color_palette
    dat[1,1] <- "one-dimensional"  
  }
  
  dat %>% 
    mutate(type = fct_relevel(type, ewm_rules)) %>%
    ggplot() +
    aes(
      x = fct_inorder(variable),
      y = pe,
      ymin = lb,
      ymax = ub,
      color = type,
      linetype = type,
      shape = type
    ) +
    geom_pointrange(
      position = position_dodge(width = 0.5) , size = 0.75) +
    scale_y_continuous(limits = limits, breaks = breaks,
                       labels = scales::number_format(accuracy = 0.01)) + 
    geom_hline(yintercept = 0) +
    geom_text(aes(
      y = y_text,
      label = paste(round(share), "%", sep = "")),
      position = position_dodge(width = 0.5),
      show.legend = FALSE 
    ) +
    scale_color_manual(values = ewm_colors) +
    theme_minimal() + 
    theme(legend.position = "bottom",
          legend.key.height = unit(1, "cm")
    ) +
    theme(panel.grid.minor = element_blank()) +
    theme(panel.grid.major.x = element_blank()) +
    theme(axis.title.x = element_text(margin = unit(c(7, 0, 0, 0), "mm"))) +
    theme(axis.title.y = element_text(margin = unit(c(0, 5, 0, 0), "mm"))) +
    labs(
      x = "Pre-Treatment Characteristics used for Targeting", 
      y = ylabel,
      color = "Treatment Rule Type",
      linetype = "Treatment Rule Type",
      shape = "Treatment Rule Type"
    )
  ggsave(paste0("./coefplot_",filename,"_textlabel",".pdf"), width = 8, height = 4.5)
}


# call function to generate and save plots
files <- c("demographics_kwh", "demographics_pmc", "demographics_smc",
           "baseline_kwh", "baseline_pmc", "baseline_smc")
map(files, make_separate_plot_text)




# COMBINED PLOTS: DEMOGRAPHICS AND PRE-TREATMENT ELECTRICITY CONSUMPTION ----

# define function to clean data (used by make_combined_plot)
clean_data <- function(dt) {
  tab.dat <- dt %>% 
    gsub("\\\\|\\(|)| ", "", .) %>% 
    read.table(text = ., sep = "&", stringsAsFactors = FALSE)
  if(tab.dat[1,1]=="ActualRCT"){
    tab.dat <- tab.dat[-(1:2),]
    tab.dat <- tab.dat[,-(4)]
  }
  dat <- cbind(tab.dat[grepl("EWM",tab.dat[,1]), ], tab.dat[!grepl("EWM",tab.dat[,1]), 4])
  names(dat) <- c("type", "variable","share","pe", "ci")
  dat <- dat %>% separate(ci, c("lb","ub"), ",")
  dat[,3:ncol(dat)] <- sapply(dat[,3:ncol(dat)], as.numeric)
  dat$type <- gsub("EWM-", "", dat$type)
  dat$type <- str_to_lower(dat$type)
  return(dat)
}

# define function to create plots
make_combined_plot_text <- function(filetype) {
  demographics.dat <- readLines(paste0("demographics_",filetype,".txt"), warn = FALSE)
  baseline.dat <- readLines(paste0("baseline_",filetype,".txt"), warn = FALSE)
  dat.list <- list(demographics.dat, baseline.dat)
  
  dat.list <- lapply(dat.list, clean_data)
  
  dat <- do.call(rbind, dat.list)
  
  dat <- dat %>% 
    mutate(variable = ifelse(type=="univariate", "only \n mean usage", variable),
           variable = ifelse(variable=="min", "min and \n mean usage", variable),
           variable = ifelse(variable=="max", "max and \n mean usage", variable),
           variable = ifelse(variable=="std", "sd and \n mean usage", variable),
           variable = ifelse(variable=="size", "house size and \n mean usage", variable),
           variable = ifelse(variable=="vintage", "house age and \n mean usage", variable),
           variable = ifelse(variable=="income", "income and \n mean usage", variable)
    )
  dat[dat$type=="univariate", "type"] <- "one-dimensional"
  
  if(filetype=="kwh"){
    ylabel <- "Electricity Consumption Changes \n (kWh per household per month)"
    limits <- c(-7,2) 
    breaks <- seq(2,-7,-2)
  } else{
    ylabel <- "Net Cost Changes \n ($ per household per month)"
    limits <- c(-1.5, 0) 
    breaks <- seq(0.5,-1.25,-0.25)
  }
  
  if (max(dat$ub)<0){ 
    y_text <- -1.5
  } else{
    y_text <- -7 
  }
  
  one_dimensional_pe <- dat[dat$type=="one-dimensional", "pe"]

  dat %>% 
    mutate(type = fct_relevel(type, c("quadrant", "cubic", "one-dimensional"))) %>%
    ggplot() +
    aes(
      x = fct_inorder(variable), 
      y = pe,
      ymin = lb,
      ymax = ub,
      color = type,
      linetype = type,
      shape = type
    ) +
    geom_pointrange(
      position = position_dodge(width = 0.75) , size = 0.75
    ) +
    scale_y_continuous(limits = limits, breaks = breaks,
                       labels = scales::number_format(accuracy = 0.01)) +
    geom_hline(yintercept = 0) +
    geom_hline(yintercept = one_dimensional_pe, color = color_palette["one-dimensional"]) +
    geom_text(aes(
      y = y_text,
      label = paste(round(share),"%",sep = "")),
      position = position_dodge(width = 0.75),
      show.legend = FALSE 
    ) +
    scale_color_manual(values = color_palette) +
    theme_minimal() + 
    theme(legend.position = "bottom"
          ,legend.key.height = unit(1, "cm")
    ) +
    theme(panel.grid.minor = element_blank()) +
    theme(panel.grid.major.x = element_blank()) + 
    theme(axis.title.x = element_text(margin = unit(c(7, 0, 0, 0), "mm"))) +
    theme(axis.title.y = element_text(margin = unit(c(0, 5, 0, 0), "mm"))) +
    labs(
      x = "Pre-Treatment Characteristics used for Targeting", 
      y = ylabel,
      color = "Treatment Rule Type",
      linetype = "Treatment Rule Type",
      shape = "Treatment Rule Type"
    )
  ggsave(paste0("./coefplot_combined_",filetype,"_textlabel",".pdf"), width = 8, height = 4.5)
}


# call function to generate and save plots
filetypes <- c("kwh", "pmc", "smc")
map(filetypes, make_combined_plot_text)
